using UnityEngine;
using System.Linq;
using System.Collections;
using System.Collections.Generic;
using DanielSig;
using DanielSig.UnityExtensions;

public class LearningAgent : MonoBehaviour, IAgent
{
	public static float ThresholdMultiplier = 1F;
	public static float StrengthIncrease = 0;
	public static float StrengthDecrease = 0;
	
	static float Sigmoid(float value)
	{
		return 2/(1+Mathf.Pow(Mathf.Epsilon, -value))-1;
	}
	static float Sigmoid01(float value)
	{
		return 1/(1+Mathf.Pow(Mathf.Epsilon, -value));
	}
	enum NeuronType
	{
		InputNeuron,
		HiddenNeuron,
		OutputNeuron
	}
	class Neuron
	{
		public Neuron(int numChannels)
		{
			Type = NeuronType.InputNeuron;
			if(numChannels > 1)
			{
				_nextChannel = new Neuron(numChannels - 1);
			}
		}
		public Neuron(NeuronType type)
		{
			Type = type;
			if(type == NeuronType.InputNeuron) _bias = 0;
		}
		public Neuron(Neuron previous, int numChannels = 0)
		{
			if(previous != null) previous.Next = this;
			Type = numChannels == 0 ? NeuronType.HiddenNeuron : NeuronType.InputNeuron;
			if(numChannels > 1)
			{
				_nextChannel = new Neuron(this, numChannels - 1);
			}
		}
		public Neuron(Neuron previous, NeuronType type)
		{
			if(previous != null) previous.Next = this;
			Type = type;
			if(type == NeuronType.InputNeuron) _bias = 0;
		}
		public float Bias
		{
			get
			{
				return _bias;
			}
			set
			{
				_bias = (_bias * Strength + value * (1-Strength)).ClampBetween(-1, 1);
			}
		}
		private float _bias = Random.value * 2 - 1;
		
		private int _numReceivedInputs = 0;
		
		public bool Firing = true;
		public NeuronType Type;
		public float Threshold = 1F;
		public float Charge = 0F;
		public float Strength = 0F;
		public LinkedList<Axiom> Inputs = new LinkedList<Axiom>();
		public LinkedList<Axiom> Outputs = new LinkedList<Axiom>();
		public Neuron Next;
		
		private float _pastCharge = 0F;
		
		private Neuron _nextChannel = null;
		public void SetCharge(params float[] charges)
		{
			Neuron n = this;
			foreach(float charge in charges)
			{
				n.Charge = charge;
				n = n._nextChannel;
				if(n == null)
					break;
			}
		}
		public void ForgetThePast(float forgetRate)
		{
			Neuron n = this;
			while(n != null)
			{
				n._pastCharge = n.Charge * forgetRate + n._pastCharge * (1-forgetRate);
				foreach(Axiom a in n.Inputs)
				{
					a.Strength = 1 - (a.Sender.Charge.GetDistance(a.Sender._pastCharge) / 2);
				}
				n = n.Next;
			}
		}
		
		public void ErrorCorrection(float correctValue, float amount)
		{	
			//divide correction evenly accross inputs
			//amount = amount / (_numReceivedInputs + 1);
			
			if(Charge != 0)//avoiding devision by zero
			{
				//float correction = amount / (_numReceivedInputs + 1);
				Bias = _bias * (1-amount) + (correctValue / Charge) * amount;
				if(correctValue.GetDistance(Charge) * amount < Threshold * ThresholdMultiplier)
					Strength += (1 - Strength) * StrengthIncrease;
				else
					Strength -= Strength * StrengthDecrease;
			}
			
			amount /= 2;
			
			foreach(Axiom a in Inputs)
			{
				if(_numReceivedInputs > 0 && !a.Sender.Firing) continue;
				//a.ErrorCorrection(correctValue, amount);
				
				if(a.Sender.Charge != 0)
				{
					//float correction = a.Sender.Charge.GetDistance(a.Sender._pastCharge) * amount;
					//half of correction is applied to this axiom
					a.Weight = a.Weight * (1-amount) + (correctValue / a.Sender.Charge) * amount;
					if(correctValue.GetDistance(a.Sender.Charge) * amount < Threshold * ThresholdMultiplier)
						a.Strength += (1 - a.Strength) * StrengthIncrease;
					else
						a.Strength -= a.Strength * StrengthDecrease;
					
				}
				//other half is applied to the rest of the hierarchy
				if(a.Sender.Type != LearningAgent.NeuronType.InputNeuron)
					a.Sender.ErrorCorrection(correctValue, amount);
			}
		}
		
		public void GatherInputs()
		{
			if(Firing) Charge = 0;
			
			Charge += Bias;
			_numReceivedInputs = 0;
			foreach(Axiom a in Inputs)
			{
				if(a.Sender.Firing)
				{
					Charge += a.Sender.Charge * a.Weight;
					_numReceivedInputs++;
				}
			}
			Firing = Type != NeuronType.OutputNeuron ? true : Charge > Threshold * ThresholdMultiplier || Charge < -Threshold * ThresholdMultiplier;
			Charge = Sigmoid(Charge);
			if(Next != null)
				Next.GatherInputs();
		}
		
		public void ConnectTo(Neuron n)
		{
			Axiom a = new Axiom{ Sender = this, Receiver = n };
			Outputs.AddLast(a);
			n.Inputs.AddLast(a);
			if(_nextChannel != null)
				_nextChannel.ConnectTo(n);
		}
		
		private Color GetColor()
		{
			Color color = Color.clear;
			Neuron n = this;
			for(int i = 0; i < 4; i++)
			{
				color[i] = n.Charge;
				n = n._nextChannel;
				if(n == null)
					break;
			}
			Color c = color.SetAlpha(1);
			c = Color.Lerp(c, Color.gray, color.a * 0.5F);
			if(n != null)
				c = c.SetAlpha(1 / n.Charge.Max(1));
			return c;
		}
		public void Draw(Vector3 position, Vector3 stepX, Vector3 stepY, float offset, bool isFirst)
		{
			Vector3 current = position + stepX * offset;
			Gizmos.color = _nextChannel != null ? GetColor() : Color.Lerp(Color.black, Color.white, Sigmoid01(Charge));//.SetAlpha(Sigmoid01(Bias).ClampBetween(0.1F, 1F));
			Gizmos.DrawSphere(current, 0.5F);
			
			position += stepY;
			offset = -Outputs.Count / 2F;
			bool first = true;
			foreach(Axiom a in Outputs)
			{
				Vector3 next = position + (stepX * offset);
				float colorGradient = Sigmoid01(Charge * a.Weight);
				Gizmos.color = Color.Lerp(Color.black, Color.white, colorGradient).SetAlpha(Sigmoid01(a.Weight).ClampBetween(0.1F, 1F));
				Gizmos.DrawLine(current, next);
				
				if(isFirst)
				{
					a.Receiver.Draw(position, stepX, stepY, offset, first);
					first = false;
				}
				offset++;
			}
		}
	}
	class Axiom
	{
		public float Weight
		{
			get
			{
				return _weight;
			}
			set
			{
				_weight = (_weight * Strength + value * (1-Strength)).ClampBetween(-1, 1);
			}
		}
		private float _weight = Random.value * 2 - 1;
		public float Strength;
		public Neuron Sender;
		public Neuron Receiver;
	}
	
	static LinkedList<Neuron> InputNeurons;
	static LinkedList<Neuron> OutputNeurons;
	static Neuron FirstHiddenNeuron;
	
	private void InitBrain (ISensor[] input)
	{
		int numInputs = input.Aggregate<ISensor,int>(0, (c,s) => c + (s.HasFixedDirection ? 1 : s.Values.Length));
		InputNeurons = new LinkedList<Neuron>();
		OutputNeurons = new LinkedList<Neuron>();
		_inputNames = new string[numInputs];
		
		//create input layer
		Neuron prevNeuron = null;
		int inputIndex = 0;
		for(int i = 0; i < numInputs; i++)
		{
			if(input[inputIndex++].HasFixedDirection)
				InputNeurons.AddLast(new Neuron(prevNeuron, 5));
			else
				InputNeurons.AddLast(new Neuron(prevNeuron, NeuronType.InputNeuron));
		}
		
		//create hidden layer
		LinkedList<Neuron> prevLayer = InputNeurons;
		LinkedList<Neuron> thisLayer = new LinkedList<Neuron>();
		for(int i = 0; i < NumberOfNueronsPerLayer.Length; i++)//for each layer
		{
			int numberOfNeurons = NumberOfNueronsPerLayer[i];
			for(int j = 0; j < numberOfNeurons; j++)//for each neuron in layer
			{
				Neuron n = prevNeuron = new Neuron(prevNeuron);
				if(FirstHiddenNeuron == null)
				{
					FirstHiddenNeuron = n;
				}
				
				thisLayer.AddLast(n);
				foreach(Neuron prev in prevLayer)
				{
					prev.ConnectTo(n);
				}
			}
			prevLayer = thisLayer;
			thisLayer = new LinkedList<Neuron>();
		}
		//create output layer
		for(int i = 0; i < 4; i++)
		{
			prevNeuron = new Neuron(prevNeuron, NeuronType.OutputNeuron);
			OutputNeurons.AddLast(prevNeuron);
		}
		
		//finally, connect hidden layer to output layer
		foreach(Neuron output in OutputNeurons)
		{
			foreach(Neuron prev in prevLayer)
			{
				prev.ConnectTo(output);
			}
		}
	}
	private void ResetOutputs ()
	{
		foreach(Neuron n in OutputNeurons)
		{
			n.Charge = 0F;
		}
	}
	private Color ColorizeSensation(SensationType type)
	{
		Color c = Color.clear;
		switch (type)
		{
			case SensationType.Projectile: c = Color.white; break;
			case SensationType.Enemy: c = Color.blue; break;
			case SensationType.EnemyAmmo: c = Color.cyan.SetAlpha(0); break;
			case SensationType.EnemyHealth: c = Color.magenta.SetAlpha(0); break;
			case SensationType.None: c = Color.black; break;
			case SensationType.Obstacle: c = Color.clear; break;
			case SensationType.AmmoPack: c = Color.green; break;
			case SensationType.MyAmmo: c = Color.green.SetAlpha(0); break;
			case SensationType.HealthPack: c = Color.red; break;
			case SensationType.MyHealth: c = Color.red.SetAlpha(0); break;
			default: return c;
		}
		c.a = 1 - c.a;
		return c;
	}
	private void RelayInputs (ISensor[] input)
	{
		LinkedListNode<Neuron> current = InputNeurons.First;
		int c = 0;
		foreach(ISensor s in input)
		{
			SensationType t = _stimulate ? ((SensationType)(int)(Random.value * 9.99999F)) : s.Type;
			if(s.HasFixedDirection)
			{
				Color color = ColorizeSensation(t);
				current.Value.SetCharge(color.r, color.g, color.b, color.a, s.Values[0]);
				current = current.Next;
				_inputNames[c++] = t.ToString();
			}
			else
			{
				int i = 0;
				foreach(float value in s.Values)
				{
					if(value == value)
						current.Value.Charge = _stimulate ? Random.value * 2 - 1 : value;
					current = current.Next;
					_inputNames[c++] = t.ToString() + (s.Values.Length > 1 ? "[" + i++ + "]" : "");
				}
			}
		}
		_stimulate = false;
		FirstHiddenNeuron.GatherInputs();
	}
	private Move GetTeacherOutput (ISensor[] input)
	{
		if(!_teacher)
		{
			_teacher = gameObject.RequestComponent<MasterAgent>();
			_teacher.StaticObjectsMask = Teacher.StaticObjectsMask;
			_teacher.OpponentTag = Teacher.OpponentTag;
			_teacher.PathFindingInterval = Teacher.PathFindingInterval;
			_teacher.ProjectilePrefabController = Teacher.ProjectilePrefabController;
		}
		
		return _teacher.NextMove(input);
	}
	private Move GetOutput ()
	{
		Move move = Move.None;
		Move current = (Move)1;
		foreach(Neuron n in OutputNeurons)
		{
			if(n.Charge >= 0)
			{
				move |= current;
			}
			current = (Move)((int)current << 1);
		}
		return move;
	}
	private void Learn(Move result, Move expected)
	{
		Move current = (Move)1;
		foreach(Neuron n in OutputNeurons)
		{
			if((result & current) != (expected & current))
				n.ErrorCorrection(-n.Charge, LearningSensitivity);
			current = (Move)((int)current << 1);
		}
	}
	
	//assigned in Unity Editor
	public LearningMethod LearningType;
	[Range(0,1)]
	public float LearningSensitivity;
	[Range(0,1)]
	public float Impulsive;
	public int[] NumberOfNueronsPerLayer;
	
	private static string[] _inputNames;
	
	private MasterAgent _teacher;
	private Move _teacherMove;
	public TeacherSettings Teacher;
	
	[System.Serializable]
	public class TeacherSettings
	{
		public LayerMask StaticObjectsMask;
		public string OpponentTag;
		public int PathFindingInterval = 20;
		public ProjectileController ProjectilePrefabController;
	}
	public enum LearningMethod
	{
		None,
		Observation,
		PunishmentAndReward,
		Both
	}
	
	public void Randomize(float amount)
	{
		float sustain = 1 - amount;
		Neuron current = FirstHiddenNeuron;
		while(current != null)
		{
			current.Strength = 0;
			current.Bias = current.Bias * sustain + (Random.value * 2 - 1) * amount;
			foreach(Axiom a in current.Inputs)
			{
				a.Strength = 0;
				a.Weight = a.Weight * sustain + (Random.value * 2 - 1) * amount;
			}
			current = current.Next;
		}
	}
	public void Stimulate()
	{
		_stimulate = true;
	}
	
	public int StandStillTimout;
	[Range(0,1)]
	public float StandStillRandomizeAmount = 0.2F;
	[Range(0,0.1F)]
	public float LearningSpeed = 0.001F;
	[Range(0,0.1F)]
	public float RelearnSpeed = 0.001F;
	[Range(0,1F)]
	public float ForgetRate = 0.1F;
	private int _standStillTimer = 0;
	private Vector3 _prevPos = Vector3.zero;
	private bool _noMove = false;
	private bool _stimulate = false;
	
	public void Start()
	{
		if(InputNeurons != null)
			InputNeurons.First.Value.ForgetThePast(1F);
	}
	
	public Move NextMove (ISensor[] input)
	{
		if(transform.position != _prevPos)
		{
			_standStillTimer = 0;
		}
		else if(_standStillTimer++ > StandStillTimout)
		{
			_standStillTimer = 0;
			Randomize(StandStillRandomizeAmount);
		}
		_prevPos = transform.position;
		
		ThresholdMultiplier = 1 - Impulsive;
		StrengthIncrease = LearningSpeed;
		StrengthDecrease = RelearnSpeed;
		
		if(InputNeurons == null)
			InitBrain(input);
		else
			ResetOutputs();
		
		InputNeurons.First.Value.ForgetThePast(ForgetRate);
		RelayInputs(input);
		
		Move neuralResponse = GetOutput();
		if(LearningType == LearningMethod.Observation || LearningType == LearningMethod.Both)
		{
			Move teacherResponse = _teacherMove = GetTeacherOutput(input);
			if(LearningSensitivity > 0) Learn(neuralResponse, teacherResponse);
		}
		_noMove = neuralResponse == Move.None || neuralResponse == (Move.TurnLeft | Move.TurnRight);
		
		return neuralResponse;
	}
	void Good(float amount)
	{
		if(OutputNeurons != null && LearningType == LearningMethod.PunishmentAndReward || LearningType == LearningMethod.Both)
		{
			if(float.IsNaN(amount)) return;
			foreach(Neuron n in OutputNeurons)
			{
				if(n.Charge != 0 || _noMove)
					n.ErrorCorrection(n.Charge.Sign() * amount, LearningSensitivity);
				//else
				//	n.ErrorCorrection(0, LearningSensitivity);
			}
		}
	}
	void Bad(float amount)
	{
		Good(-amount);
	}
	void OnDrawGizmos()
	{
		if(!Application.isPlaying) return;
		
		Vector3 origin = Camera.main.ViewportToWorldPoint(DrawOrigin);
		bool first = true;
		float offset = -InputNeurons.Count / 2F;
		int c = 0;
		foreach(Neuron n in InputNeurons)
		{
			
			(_inputNames[c++] ?? "null").LogAt(origin + DrawStepX * offset - DrawStepY);
			n.Draw(origin, DrawStepX, DrawStepY, offset++, first);
			first = false;
		}
		
		origin += DrawStepY * (NumberOfNueronsPerLayer.Length + 1.2F);
		offset = -2;
		Move current = (Move)1;
		for(int i = 0; i < 4; i++)
		{
			Gizmos.color = (int)(_teacherMove & current) > 0 ? Color.green : Color.red;
			Gizmos.DrawSphere(origin + DrawStepX * offset++, 0.5F);
			
			current = (Move)((int)current << 1);
		}
	}
	public Vector3 DrawStepX = Vector3.right;
	public Vector3 DrawStepY = Vector3.forward;
	public Vector3 DrawOrigin = Vector3.forward;
}

